from typing import Annotated

from langchain_core.tools import tool
from pydantic import BaseModel, Field


class CalculateArgs(BaseModel):
    a: float = Field(description="第一个数字")
    b: float = Field(description="第二个数字")
    operation: str = Field(description="运算符，必须是+ - * /")


@tool('calculate')
def calculate(
        a: Annotated[float, Field(description="第一个数字")],
        b: Annotated[float, Field(description="第二个数字")],
        operation: Annotated[str, Field(description="运算符，必须是+ - * /")]) -> float:
    """
        工具函数：计算两个数字的运算结果
        参数：
        a：浮点型，第一个数字
        b：浮点型，第二个参数
        operation：字符串，运算符
    """
    result = 0.0
    match operation:
        case '+':
            result = a + b
        case '-':
            result = a - b
        case '*':
            result = a * b
        case '/':
            if b == 0:
                raise ZeroDivisionError("除数不能为0")
            result = a / b
        case _:
            result = 0.0

    return result


# print(f'工具名称：{calculate.name}')
# print(f'工具描述：{calculate.description}')
# print(f'工具参数：{calculate.args}')
# print(f'是否直接返回：{calculate.return_direct}')
# print(f'工具schema：{calculate.args_schema.model_json_schema()}')

res = calculate.invoke({"a": 1, "b": 2, "operation": "+"})
print(res)
